#!/usr/bin/python3
# -*- coding: utf-8 -*-
import datetime
import json
import os

from pycocotools.coco import COCO

CLASSES = ['person', 'cat', 'dog', 'baby_carriage', 'face', 'cell phone', 'bicycle', 'motorcycle']
# CLASSES = ['person']
# CLASSES = ['cat']
# CLASSES = ['dog']
# CLASSES = ['cell phone']
# CLASSES = ['bicycle']
# CLASSES = ['motorcycle']

# SKIP_CROWD = True
SKIP_CROWD = False

ANNO_FOLDER = r'/root/data/coco/annotations/'
ANNO_FILENAME = r'instances_val2017.json'
OUT_ANNO_FILENAME = r'instances_select_val2017.json'
# OUT_ANNO_FILENAME = r'person_instances_select_val2017.json'
# OUT_ANNO_FILENAME = r'cat_instances_select_val2017.json'
# OUT_ANNO_FILENAME = r'dog_instances_select_val2017.json'
# OUT_ANNO_FILENAME = r'phone_dog_instances_select_val2017.json'
# OUT_ANNO_FILENAME = r'bicycle_dog_instances_select_val2017.json'
# OUT_ANNO_FILENAME = r'motorcycle_dog_instances_select_val2017.json'


# ANNO_FOLDER = r'/root/data/coco/annotations/'
# ANNO_FILENAME = r'instances_train2017.json'
# OUT_ANNO_FILENAME = r'instances_select_train2017.json'

# Val
# person 11004
# bicycle 316
# motorcycle 371
# cat 202
# dog 218
# cell phone 262

# Train
# person 262465
# bicycle 7113
# motorcycle 8725
# cat 4768
# dog 5508
# cell phone 6434


INFO = {
    "description": "Erised Dataset",
    "url": "http://www.erised.com",
    "version": "0.1.0",
    "year": 2020,
    "contributor": "Erised",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
        "id": 1,
        "name": "Attribution-NonCommercial-ShareAlike License",
        "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/"
    }
]

CATEGORIES = [
    {
        'id': 1,
        'name': 'person',
        'supercategory': 'person',  # person
    },
    {
        'id': 17,
        'name': 'cat',
        'supercategory': 'cat',  # animal
    },
    {
        'id': 18,
        'name': 'dog',
        'supercategory': 'dog',  # animal
    },
    {
        'id': 102,
        'name': 'baby_carriage',
        'supercategory': 'baby_carriage',
    },
    {
        'id': 103,
        'name': 'face',
        'supercategory': 'face',  # person
    },
    {
        'id': 77,
        'name': 'cell phone',
        'supercategory': 'cell phone',  # electronic
    },
    {
        'id': 2,
        'name': 'bicycle',
        'supercategory': 'bicycle',  # vehicle
    },
    {
        'id': 4,
        'name': 'motorcycle',
        'supercategory': 'motorcycle',  # vehicle
    },
]

# CATEGORIES = [
#     {"supercategory": "person","id": 1,"name": "person"},{"supercategory": "vehicle","id": 2,"name": "bicycle"},{"supercategory": "vehicle","id": 3,"name": "car"},{"supercategory": "vehicle","id": 4,"name": "motorcycle"},{"supercategory": "vehicle","id": 5,"name": "airplane"},{"supercategory": "vehicle","id": 6,"name": "bus"},{"supercategory": "vehicle","id": 7,"name": "train"},{"supercategory": "vehicle","id": 8,"name": "truck"},{"supercategory": "vehicle","id": 9,"name": "boat"},{"supercategory": "outdoor","id": 10,"name": "traffic light"},{"supercategory": "outdoor","id": 11,"name": "fire hydrant"},{"supercategory": "outdoor","id": 13,"name": "stop sign"},{"supercategory": "outdoor","id": 14,"name": "parking meter"},{"supercategory": "outdoor","id": 15,"name": "bench"},{"supercategory": "animal","id": 16,"name": "bird"},{"supercategory": "animal","id": 17,"name": "cat"},{"supercategory": "animal","id": 18,"name": "dog"},{"supercategory": "animal","id": 19,"name": "horse"},{"supercategory": "animal","id": 20,"name": "sheep"},{"supercategory": "animal","id": 21,"name": "cow"},{"supercategory": "animal","id": 22,"name": "elephant"},{"supercategory": "animal","id": 23,"name": "bear"},{"supercategory": "animal","id": 24,"name": "zebra"},{"supercategory": "animal","id": 25,"name": "giraffe"},{"supercategory": "accessory","id": 27,"name": "backpack"},{"supercategory": "accessory","id": 28,"name": "umbrella"},{"supercategory": "accessory","id": 31,"name": "handbag"},{"supercategory": "accessory","id": 32,"name": "tie"},{"supercategory": "accessory","id": 33,"name": "suitcase"},{"supercategory": "sports","id": 34,"name": "frisbee"},{"supercategory": "sports","id": 35,"name": "skis"},{"supercategory": "sports","id": 36,"name": "snowboard"},{"supercategory": "sports","id": 37,"name": "sports ball"},{"supercategory": "sports","id": 38,"name": "kite"},{"supercategory": "sports","id": 39,"name": "baseball bat"},{"supercategory": "sports","id": 40,"name": "baseball glove"},{"supercategory": "sports","id": 41,"name": "skateboard"},{"supercategory": "sports","id": 42,"name": "surfboard"},{"supercategory": "sports","id": 43,"name": "tennis racket"},{"supercategory": "kitchen","id": 44,"name": "bottle"},{"supercategory": "kitchen","id": 46,"name": "wine glass"},{"supercategory": "kitchen","id": 47,"name": "cup"},{"supercategory": "kitchen","id": 48,"name": "fork"},{"supercategory": "kitchen","id": 49,"name": "knife"},{"supercategory": "kitchen","id": 50,"name": "spoon"},{"supercategory": "kitchen","id": 51,"name": "bowl"},{"supercategory": "food","id": 52,"name": "banana"},{"supercategory": "food","id": 53,"name": "apple"},{"supercategory": "food","id": 54,"name": "sandwich"},{"supercategory": "food","id": 55,"name": "orange"},{"supercategory": "food","id": 56,"name": "broccoli"},{"supercategory": "food","id": 57,"name": "carrot"},{"supercategory": "food","id": 58,"name": "hot dog"},{"supercategory": "food","id": 59,"name": "pizza"},{"supercategory": "food","id": 60,"name": "donut"},{"supercategory": "food","id": 61,"name": "cake"},{"supercategory": "furniture","id": 62,"name": "chair"},{"supercategory": "furniture","id": 63,"name": "couch"},{"supercategory": "furniture","id": 64,"name": "potted plant"},{"supercategory": "furniture","id": 65,"name": "bed"},{"supercategory": "furniture","id": 67,"name": "dining table"},{"supercategory": "furniture","id": 70,"name": "toilet"},{"supercategory": "electronic","id": 72,"name": "tv"},{"supercategory": "electronic","id": 73,"name": "laptop"},{"supercategory": "electronic","id": 74,"name": "mouse"},{"supercategory": "electronic","id": 75,"name": "remote"},{"supercategory": "electronic","id": 76,"name": "keyboard"},{"supercategory": "electronic","id": 77,"name": "cell phone"},{"supercategory": "appliance","id": 78,"name": "microwave"},{"supercategory": "appliance","id": 79,"name": "oven"},{"supercategory": "appliance","id": 80,"name": "toaster"},{"supercategory": "appliance","id": 81,"name": "sink"},{"supercategory": "appliance","id": 82,"name": "refrigerator"},{"supercategory": "indoor","id": 84,"name": "book"},{"supercategory": "indoor","id": 85,"name": "clock"},{"supercategory": "indoor","id": 86,"name": "vase"},{"supercategory": "indoor","id": 87,"name": "scissors"},{"supercategory": "indoor","id": 88,"name": "teddy bear"},{"supercategory": "indoor","id": 89,"name": "hair drier"},{"supercategory": "indoor","id": 90,"name": "toothbrush"}
# ]


def coco_select_classes_data(anno_folder, anno_filename, classes_list, out_anno_filename):
    coco_output = {
        "info": INFO,
        "licenses": LICENSES,
        "categories": CATEGORIES,
        "images": [],
        "annotations": []
    }

    anno_file = os.path.join(anno_folder, anno_filename)
    coco = COCO(anno_file)

    # get all images containing given categories, select one at random
    catIds = coco.getCatIds(catNms=classes_list)
    # catIds = coco.getCatIds()
    _img_ids = set()
    for cat_id in catIds:
        _img_ids |= set(coco.getImgIds(catIds=cat_id))
    img_ids = list(_img_ids)
    # imgIds = coco.getImgIds()

    objCntPreCls = {}
    for id in catIds:
        objCntPreCls[id] = 0

    file_count = 0
    for i in img_ids:
        file_count = file_count + 1
        if (file_count % 100) == 0:
            print('Proccess', file_count, 'Files.')

        image_info = coco.loadImgs([i])[0]
        image_info["license"] = 1
        image_info["coco_url"] = ""
        image_info["flickr_url"] = ""

        img_id = image_info['id']
        ann_ids = coco.getAnnIds(imgIds=[img_id])
        ann_info = coco.loadAnns(ann_ids)
        ann_list = []
        for j, ann in enumerate(ann_info):
            if ann["category_id"] not in catIds:
                continue
            # 不要标注为成群的图片，只要一个一个单独标注的图片
            if SKIP_CROWD:
                if ann["iscrowd"]:
                    ann_list.clear()
                    break
            ann["segmentation"] = []
            # "area": area
            objCntPreCls[ann["category_id"]] += 1
            ann_list.append(ann)

        if len(ann_list) > 0:
            coco_output["images"].append(image_info)
            for ann in ann_list:
                coco_output["annotations"].append(ann)

    print('Proccess', file_count, 'Files.')

    for id, objCnt in objCntPreCls.items():
        for c in CATEGORIES:
            if c['id'] == id:
                print(c['name'], objCnt)
                break

    out_json_filepath = os.path.join(anno_folder, out_anno_filename)
    with open(out_json_filepath, 'w') as output_json_file:
        json.dump(coco_output, output_json_file)
    print('Save ', out_json_filepath, 'Finish!')


if __name__ == "__main__":
    coco_select_classes_data(ANNO_FOLDER, ANNO_FILENAME, CLASSES, OUT_ANNO_FILENAME)
